package zuo.algo.c08;

import zuo.algo.c12.UnionFind;
import zuo.algo.data.graph.Edge;
import zuo.algo.data.graph.Graph;
import zuo.algo.data.graph.Node;
import zuo.algo.util.GraphUtils;

import java.util.*;

/**
 * 01:12-Kruskal(最小生成树)算法
 * 01:41-Prim(最小生成树)算法
 */
public class MinimumSpanningTree {

    public static List<Edge> k(Graph graph) {
        PriorityQueue<Edge> priorityQueue = new PriorityQueue<>((v1, v2) -> {
            if (v1.weight < v2.weight) {
                return -1;
            } else if (v1.weight > v2.weight) {
                return 1;
            }
            return 0;
        });
        priorityQueue.addAll(graph.edges);
        UnionFind.UnionFindSet<Node> unionFind = new UnionFind
                .UnionFindSet<>(new ArrayList<>(graph.nodes.values()));
        Set<Edge> res = new HashSet<>();
        while (!priorityQueue.isEmpty()) {
            Edge edge = priorityQueue.poll();
            if (!unionFind.isSameSet(edge.from, edge.to)) {
                unionFind.union(edge.from, edge.to);
                res.add(edge);
            }
        }
        return new ArrayList<>(res);
    }

    public static List<Edge> p(Node node) {
        List<Edge> res = new ArrayList<>();

        PriorityQueue<Edge> priorityQueue = new PriorityQueue<>((v1, v2) -> {
            if (v1.weight < v2.weight) {
                return -1;
            } else if (v1.weight > v2.weight) {
                return 1;
            }
            return 0;
        });

        Set<Node> usedNodeSet = new HashSet<>();
        usedNodeSet.add(node);
        addEdgeToPq(priorityQueue, usedNodeSet, node);

        while (!priorityQueue.isEmpty()) {
            Edge edge = priorityQueue.poll();
            node = edge.to;
            if (!usedNodeSet.contains(node)) {
                res.add(edge);
                usedNodeSet.add(node);
                addEdgeToPq(priorityQueue, usedNodeSet, node);
            }
        }
        return res;
    }

    private static void addEdgeToPq(PriorityQueue<Edge> pq, Set<Node> usedNodeSet, Node node) {
        ArrayList<Edge> edges = node.edges;
        for (Edge edge : edges) {
            if (!usedNodeSet.contains(edge.to)) {
                pq.add(edge);
            }
        }
    }

    public static void main(String[] args) {
        int[][] matrix = {
                {0, 2, 20, 13, 0, 0, 0, 0},
                {2, 0, 4, 0, 0, 0, 0, 0},
                {20, 4, 0, 3, 100, 10, 0, 0},
                {13, 0, 3, 0, 0, 0, 0, 0},
                {0, 0, 100, 0, 0, 0, 0, 0},
                {0, 0, 10, 0, 0, 0, 9, 1},
                {0, 0, 0, 0, 0, 9, 0, 0},
                {0, 0, 0, 0, 0, 1, 0, 0},
        };
        Graph graph = GraphUtils.createGraph1(matrix);
        List<Edge> p = p(graph.nodes.get(0));
        List<Edge> k = k(graph);
        System.out.println(p);
    }
}
